import aubio
import numpy as np
import pretty_midi
import soundfile as sf
import logging



logging.basicConfig(
    level=logging.INFO
)
logger = logging.getLogger(__name__)

import numpy as np
import pathlib 
def audio_to_midi(audio_file:str|pathlib.Path)->str:
    logger.info(f"Received file {audio_file}")
    # 确保是 float32
    audio, sr = sf.read(audio_file, dtype='float32')
    if audio.ndim > 1:
        audio = np.mean(audio, axis=1)  # 转单声道
    audio = audio.astype(np.float32)
    logger.info(f"buffer: {audio[:10]}")

    # aubio pitch
    win_s = 2048
    hop_s = 512
    pitch_o = aubio.pitch("yin", win_s, hop_s, sr)
    pitch_o.set_unit("Hz")
    pitch_o.set_silence(-40)

    midi = pretty_midi.PrettyMIDI()
    instrument = pretty_midi.Instrument(program=0)
    logger.info(f"midi: {midi}")
    logger.info(f"instrument: {instrument}")

    for i in range(0, len(audio) - hop_s, hop_s):
        print(f"processing page No. {i} ({100.0*i/len(audio):2.2f}%) ...",end='\r')
        frame = audio[i:i+hop_s].astype(np.float32)  # 用 hop_s 而不是 win_s
        if len(frame) < hop_s:
            continue
        f0 = pitch_o(frame)[0]

        if f0 > 0:
            midi_note = int(round(pretty_midi.hz_to_note_number(f0)))
            midi_note = np.clip(midi_note, 0, 127)  # clamp
            note = pretty_midi.Note(
                velocity=100, pitch=midi_note,
                start=i/float(sr), end=(i+win_s)/float(sr)
            )
            instrument.notes.append(note)

    midi.instruments.append(instrument)
    midi_file = audio_file + ".mid"
    midi.write(midi_file)
    logger.info(f"Wrote {midi_file}")

    return midi_file

def show_output(midi_file:str|pathlib.Path):
    import pypianoroll
    import matplotlib.pyplot as plt
    import matplotlib
    matplotlib.use("QtAgg")  # 或者 "Qt5Agg"
    multitrack = pypianoroll.read(midi_file)
    multitrack.plot()
    plt.show()


if __name__=="__main__":
    import os
    from soundfile import  available_formats
    candidates = os.listdir(".")
    
    input_file = [f for f in os.listdir(".") if f.split(".")[-1].upper() in available_formats().keys()][0]
    # 使用示例
    # output = audio_to_midi(input_file)
    midi_file = input_file + ".mid"
    show_output(midi_file)